import gymnasium as gym
import gymnasium as gym.spaces
from gym.utils import seeding
import enum
import numpy as np

from . import data

DEFAULT_BARS_COUNT = 10
DEFAULT_COMMISSION_PERC = 0.1


# 购买股票的动作
class Actions(enum.Enum):
    # 不采取行动
    Skip = 0
    # 购买单个股票
    Buy = 1
    # 平仓
    Close = 2


class State:
    '''
    股票环境状态，表示当前股票环境的的情况
    '''

    def __init__(self, bars_count, commission_perc, reset_on_close, reward_on_close=True, volumes=True):
        '''
        构造函数仅仅只是验证输入的数据类型是否正确，然后保存起来
        bars_count：每次进行环境观察时，数据采样的大小
        commission_perc: 每次买卖股票是需要支付的佣金
        reset_on_close: 每次平仓后，是重新开始新的一轮训练还是继续训练一直到时间结束
        reward_on_close：是否每次平仓后都计算更新回报
        volumes: 是否在观察数据中记录成交量
        '''
        assert isinstance(bars_count, int)
        assert bars_count > 0
        assert isinstance(commission_perc, float)
        assert commission_perc >= 0.0
        assert isinstance(reset_on_close, bool)
        assert isinstance(reward_on_close, bool)
        self.bars_count = bars_count
        self.commission_perc = commission_perc
        self.reset_on_close = reset_on_close
        self.reward_on_close = reward_on_close
        self.volumes = volumes

    def reset(self, prices, offset):
        '''
        todo 补充含义
        重置环境状态，重置为没有购买股票的状态
        '''
        assert isinstance(prices, data.Prices)
        assert offset >= self.bars_count-1
        # have_position: 是否已经购买了股票
        self.have_position = False
        # todo 购买价格设置为0
        self.open_price = 0.0
        self._prices = prices
        self._offset = offset

    @property
    def shape(self):
        # [h, l, c] * bars + position_flag + rel_profit (since open)
        # 这个shape值得是观察的股票购买量空间长度 + 已购买股票份额 + 基于开盘价的收益率
        # 有bars个[h（最高价）,l（最低价）,c（收盘价）]，每个[h,l,c]有3个值，所以有3*bars个值
        # 根据股票价格的存储形式来获取状态的返回shape
        # 对于在观察中启用交易量self.volumes，每个[h,l,c,v]有4个值，所以有4*bars个值
        if self.volumes:
            return (4 * self.bars_count + 1 + 1, )
        else:
            return (3*self.bars_count + 1 + 1, )

    def encode(self):
        """
        Convert current state into numpy array.
        转换当前的状态数据到numpy数组类型
        每个的状态都包含观察空间长度的交易量数据，以及已购买的股票份额，基于开盘价的收益率
        # todo 补充转换过程
        """
        res = np.ndarray(shape=self.shape, dtype=np.float32)
        shift = 0
        for bar_idx in range(-self.bars_count+1, 1):
            # 根据state的存储形式将股票数据按照[高价、低价、收盘价，if self.volumnes 买入买出的量]
            res[shift] = self._prices.high[self._offset + bar_idx]
            shift += 1
            res[shift] = self._prices.low[self._offset + bar_idx]
            shift += 1
            res[shift] = self._prices.close[self._offset + bar_idx]
            shift += 1
            # todo 购买量的作用，因为这里没有看出回报和购买量之间的关系
            if self.volumes:
                res[shift] = self._prices.volume[self._offset + bar_idx]
                shift += 1
        # todo 这句的含义是什么？
        res[shift] = float(self.have_position)
        shift += 1
        # todo 下面这些的含义
        if not self.have_position:
            # 如果没有购买，则增长率为0
            res[shift] = 0.0
        else:
            # 如果已经购买，则这个为基于购买价的利润
            res[shift] = (self._cur_close() - self.open_price) / self.open_price
        return res

    def _cur_close(self):
        """
        Calculate real close price for the current bar
        对于当前条目的价格，计算真实的收盘价，收盘价是根据开盘价为基准进行的偏移价格
        在prices中，收盘价是与开盘价的相对比率，所以这边需要计算出真实的收盘价
        """
        open = self._prices.open[self._offset]
        rel_close = self._prices.close[self._offset]
        return open * (1.0 + rel_close)

    def step(self, action):
        """
        Perform one step in our price, adjust offset, check for the end of prices
        and handle position change
        根据当前的状态，执行一次动作，获得执行动作完成后的状态
        :param action:
        :return: reward, done，本次指定动作获取的回报以及是否已经结束
        """
        assert isinstance(action, Actions)
        reward = 0.0
        done = False
        close = self._cur_close()
        # have_position是表示是否已经购买了股票
        if action == Actions.Buy and not self.have_position:
            # 如果当前执行的动作是买，并且上次并没有执行买动作，则将
            # 存储动作标记并且计算此时购买的价格作为开盘价参考
            # 该模型中，假设在购买时都是以收盘价购买的
            self.have_position = True
            self.open_price = close
            # 每次购买都需要支付佣金，所以回报需要减去佣金
            reward -= self.commission_perc
        # 如果执行的动作是卖，并且已经购买了股票
        elif action == Actions.Close and self.have_position:
            # 每次卖出都需要支付佣金，所以回报需要减去佣金
            reward -= self.commission_perc
            # 根据reset_on_close来判定本次执行动作是否结束还是等待所有的数据训练完成后结束
            done |= self.reset_on_close
            if self.reward_on_close:
                # 根据卖出的价格计算回报
                # 这里之所要乘以100，是因为这里的reward的金额都是以百分比的形式表示的
                # 这里的计算结果表示：卖出的价格比开盘价高多少百分比，也就是收益多少百分比
                reward += 100.0 * (close - self.open_price) / self.open_price
            self.have_position = False
            self.open_price = 0.0

        # 得到下一个股票数据
        self._offset += 1
        # 记录本次的收盘价格
        prev_close = close
        # 下一次的收盘价格，因为self.offset已经移动
        close = self._cur_close()
        # 判断训练数据是否已经结束
        done |= self._offset >= self._prices.close.shape[0]-1

        # 如果当前是已经购买股票，并且计算激励是每次训练都计算一次回报
        # 则根据当前的收盘价以及上一次的收盘价计算本次的回报激励
        # 如果已经购买股票，且设置的是每次执行完都更新回报（而不是等卖出才计算回报）
        # 则这里就直接更新回报
        if self.have_position and not self.reward_on_close:
            reward += 100.0 * (close - prev_close) / prev_close

        # 返回基于百分比的回报以及是否结束
        return reward, done


class State1D(State):
    """
    State with shape suitable for 1D convolution
    """
    @property
    def shape(self):
        # shape与State不同
        # shape 表示 最高价、最低价、收盘价、if self.volumnes 买入买出的量、观察的样本长度
        if self.volumes:
            return (6, self.bars_count)
        else:
            return (5, self.bars_count)

    def encode(self):
        '''
        将当前状态编码为一个向量
        '''

        res = np.zeros(shape=self.shape, dtype=np.float32)
        ofs = self.bars_count-1
        # 从价格样本中采样ofs长度的最高价、最低价、收盘价
        res[0] = self._prices.high[self._offset-ofs:self._offset+1]
        res[1] = self._prices.low[self._offset-ofs:self._offset+1]
        res[2] = self._prices.close[self._offset-ofs:self._offset+1]
        if self.volumes:
            # 如果有需要统计买入卖出量则需要将其存储起来
            res[3] = self._prices.volume[self._offset-ofs:self._offset+1]
            dst = 4
        else:
            dst = 3
        if self.have_position:
            # 如果有购买股票，则res[dst]表示购买标识
            # res[dst+1]表示相对于购买价的增长率 ，也就是基于开盘价的收益率，也就是回报
            res[dst] = 1.0
            res[dst+1] = (self._cur_close() - self.open_price) / self.open_price
        return res


# 构建股票交易环境
class StocksEnv(gym.Env):
    # 这个仅作兼容才实现的元数据，并未有任何实际意义，可以忽略
    metadata = {'render.modes': ['human']}

    def __init__(self, prices, bars_count=DEFAULT_BARS_COUNT,
                 commission=DEFAULT_COMMISSION_PERC, reset_on_close=True, state_1d=False,
                 random_ofs_on_reset=True, reward_on_close=False, volumes=False):
        '''
        prices: 字典类型，包含股票的价格数据，也就是可以包含多个文件的股票训练数据
        state_1d: 决定训练数据编码的排序顺序，具体查看书籍P160页
        bars_count: 决定每次观察的股票数据量，也决定了每次观察的状态向量的长度
        commission：买卖股票时必须支付给经纪人的佣金比例
        reset_on_close： true,表示每次平仓后则重置本次训练集的参数；否则在训练数据时间序列结束后重置，todo 重置的参数有哪些
        '''
        assert isinstance(prices, dict)
        self._prices = prices
        if state_1d:
            # todo 理解State1D的代码
            self._state = State1D(bars_count, commission, reset_on_close, reward_on_close=reward_on_close,
                                  volumes=volumes)
        else:
            # state_1d为false时，State将股票的数据以[最高价、最低价、收盘价、购买量、最高价...]的顺序编码以一维向量的方式进行编码
            self._state = State(bars_count, commission, reset_on_close, reward_on_close=reward_on_close,
                                volumes=volumes)
        # 通过gym创建离线动作空间
        # todo 补充Discrete的说明
        self.action_space = gym.spaces.Discrete(n=len(Actions))
        # 通过gym创建观察空间，得到状态数据
        # todo 补充observation_space的说明
        # todo 为什么self._state.shape的观察空间是32维度的
        self.observation_space = gym.spaces.Box(low=-np.inf, high=np.inf, shape=self._state.shape, dtype=np.float32)
        self.random_ofs_on_reset = random_ofs_on_reset
        self.seed()

    def reset(self):
        # make selection of the instrument and it's offset. Then reset the state
        # 从多个股票训练数据csv文件中随机选择一个数据集，开始进行训练
        self._instrument = self.np_random.choice(list(self._prices.keys()))
        prices = self._prices[self._instrument]
        bars = self._state.bars_count
        # 判断是否需要随机获取一个状态起始位置
        # 否则就从其实位置开始
        if self.random_ofs_on_reset:
            offset = self.np_random.choice(prices.high.shape[0]-bars*10) + bars
        else:
            offset = bars
        # 根据选择的股票类型以及偏移位置，重置状态
        self._state.reset(prices, offset)
        # todo 了解encode函数的作用,书上说是用来请求初始状态，根据step中的代码来看
        #   是用来获取当前状态空间的状态数据
        return self._state.encode()

    def step(self, action_idx):
        '''
        执行一次动作
        action_idx: 动作的id
        return
        obs:下一轮观察的数据
        reward: 执行本次动作后所获取的激励
        done：是否结束
        info: todo 未知补充 从代码中可知，其中包含了本次起始状态下的股票名称和偏移位置
        '''
        action = Actions(action_idx)
        # todo 了解step类的作用
        reward, done = self._state.step(action)
        # 每次执行完step后，已经填入到numpy中的观察数据已经被使用过了
        # 所以需要重新读取股票数据到numpy缓存中
        # 每次执行一步后，都会将当前状态及前bar_count个状态返回，作为观察状态数据
        obs = self._state.encode()
        info = {"instrument": self._instrument, "offset": self._state._offset}
        return obs, reward, done, info

    def render(self, mode='human', close=False):
        pass

    def close(self):
        pass

    def seed(self, seed=None):
        # todo 补充随机数种子如何实现
        self.np_random, seed1 = seeding.np_random(seed)
        seed2 = seeding.hash_seed(seed1 + 1) % 2 ** 31
        return [seed1, seed2]

    @classmethod
    def from_dir(cls, data_dir, **kwargs):
        '''
        从文件夹中读取股票数据文件进行构建
        一般类方法用于构建类
        '''
        # todo 了解load_relative方法
        prices = {file: data.load_relative(file) for file in data.price_files(data_dir)}
        return StocksEnv(prices, **kwargs)
